import time
import psutil
import subprocess
import json
import os
import tkinter as tk
from tkinter import messagebox, filedialog
import winreg
import sys
from PIL import Image
import pystray
import threading

# 全局变量用于存储系统托盘图标对象
icon = None

def get_exe_name(path):
    if path.find('"') != -1:
        path = path.replace('"', '')
    path = path.replace('\\', '/')
    return path.split('/')[-1]

def load_config():
    if getattr(sys, 'frozen', False):
        # 如果是打包后的 exe 文件
        program_dir = os.path.dirname(sys.executable)
    else:
        # 如果是运行 Python 脚本
        program_dir = os.path.dirname(os.path.abspath(__file__))
    config_path = os.path.join(program_dir, 'config.json')
    if os.path.exists(config_path):
        try:
            with open(config_path, 'r', encoding='utf-8') as f:
                config = json.load(f)
                if isinstance(config.get("processes"), list):
                    return config
                else:
                    print("Invalid processes format in config file, expected a list")
                    return {"processes": []}
        except json.JSONDecodeError as e:
            print(f"Invalid JSON format in config file: {str(e)}")
            return {"processes": []}
    else:
        return {"processes": []}

def save_config(config):
    if getattr(sys, 'frozen', False):
        # 如果是打包后的 exe 文件
        program_dir = os.path.dirname(sys.executable)
    else:
        # 如果是运行 Python 脚本
        program_dir = os.path.dirname(os.path.abspath(__file__))
    config_path = os.path.join(program_dir, 'config.json')
    try:
        with open(config_path, 'w', encoding='utf-8') as f:
            json.dump(config, f, indent=4)
    except Exception as e:
        messagebox.showerror("错误", f"保存配置文件失败: {str(e)}")

def start_pid_alive():
    config = load_config()
    process_list = config.get("processes", [])
    
    for process in process_list:
        path = process.get("path", "")
        config_path = process.get("config", "")
        exe_name = get_exe_name(path)
        
        process_running = False
        for proc in psutil.process_iter():
            try:
                if proc.name() == exe_name:
                    process_running = True
                    break
            except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
                continue

        if not process_running:
            try:
                args = [path]
                if config_path:
                    if not os.path.isabs(config_path):
                        config_path = os.path.join(os.path.dirname(path), config_path)
                    args.extend(["--config", config_path])
                subprocess.Popen(args, cwd=os.path.dirname(path))
                print(f"Process {exe_name} started with args: {args[1:]}")
            except Exception as e:
                print(f"Failed to start process {path}: {str(e)}")
        else:
            print(f"Process {exe_name} is already running.")

def monitor_processes():
    while True:
        start_pid_alive()
        time.sleep(10)

def set_startup():
    try:
        key = winreg.OpenKey(winreg.HKEY_CURRENT_USER,
                             r"Software\Microsoft\Windows\CurrentVersion\Run",
                             0, winreg.KEY_SET_VALUE)
        exe_path = f'"{os.path.abspath(__file__)}" minimize'
        winreg.SetValueEx(key, "ProcessGuard", 0, winreg.REG_SZ, exe_path)
        winreg.CloseKey(key)
        messagebox.showinfo("成功", "已设置开机自启")
    except Exception as e:
        messagebox.showerror("错误", f"设置开机自启失败: {str(e)}")

def remove_startup():
    try:
        key = winreg.OpenKey(winreg.HKEY_CURRENT_USER,
                             r"Software\Microsoft\Windows\CurrentVersion\Run",
                             0, winreg.KEY_ALL_ACCESS)
        winreg.DeleteValue(key, "ProcessGuard")
        winreg.CloseKey(key)
        messagebox.showinfo("成功", "已移除开机自启")
    except Exception as e:
        messagebox.showerror("错误", f"移除开机自启失败: {str(e)}")

def select_program_path():
    file_path = filedialog.askopenfilename(filetypes=[("可执行文件", "*.exe")])
    if file_path:
        program_path_entry.delete(0, tk.END)
        program_path_entry.insert(0, file_path)

def select_config_path():
    file_path = filedialog.askopenfilename(filetypes=[("JSON 文件", "*.json")])
    if file_path:
        config_path_entry.delete(0, tk.END)
        config_path_entry.insert(0, file_path)

def add_process():
    program_path = program_path_entry.get()
    config_path = config_path_entry.get()

    if program_path:
        config = load_config()
        new_process = {
            "path": program_path,
            "config": config_path
        }
        config["processes"].append(new_process)
        save_config(config)
        messagebox.showinfo("成功", "已添加新的监控进程")

def show_window():
    global icon
    if icon:
        icon.stop()
    root.deiconify()

def quit_app(icon, item):
    icon.stop()
    root.destroy()
    os._exit(0)

def hide_window():
    global icon
    root.withdraw()
    icon_path = 'icon.png'  # 请替换为你的图标文件路径
    try:
        image = Image.open(icon_path)
        menu = pystray.Menu(
            pystray.MenuItem('显示窗口', show_window),
            pystray.MenuItem('退出', quit_app)
        )
        icon = pystray.Icon("ProcessGuard", image, "进程守护程序", menu)
        icon.run()
    except FileNotFoundError:
        error_msg = f"未找到系统托盘图标文件，请检查路径: {os.path.abspath(icon_path)}"
        messagebox.showerror("错误", error_msg)

root = tk.Tk()
root.title("进程守护程序")

# 程序路径选择
program_path_label = tk.Label(root, text="程序路径:")
program_path_label.pack(pady=5)
program_path_entry = tk.Entry(root, width=50)
program_path_entry.pack(pady=5)
program_path_button = tk.Button(root, text="选择程序路径", command=select_program_path)
program_path_button.pack(pady=5)

# 配置文件路径选择
config_path_label = tk.Label(root, text="配置文件路径:")
config_path_label.pack(pady=5)
config_path_entry = tk.Entry(root, width=50)
config_path_entry.pack(pady=5)
config_path_button = tk.Button(root, text="选择配置文件路径", command=select_config_path)
config_path_button.pack(pady=5)

# 添加进程按钮
add_button = tk.Button(root, text="添加进程", command=add_process)
add_button.pack(pady=20)

startup_button = tk.Button(root, text="设置开机自启", command=set_startup)
startup_button.pack(pady=10)

remove_startup_button = tk.Button(root, text="移除开机自启", command=remove_startup)
remove_startup_button.pack(pady=10)

# 启动进程监控线程
monitor_thread = threading.Thread(target=monitor_processes, daemon=True)
monitor_thread.start()

# 默认最小图标运行
root.withdraw()
root.after(0, hide_window)

root.protocol('WM_DELETE_WINDOW', hide_window)

root.mainloop()
